#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
@Time : 2022/4/15
@Author : 潘世友
"""

from datetime import datetime
from typing import Literal

from dateutil.relativedelta import relativedelta

TODAY = datetime.today().strftime('%Y-%m-%d')


def _year(indate: datetime, sign: Literal['s', 'b', 'e'] = 's') -> datetime:
    if sign.lower() in ['s', 'same']:
        outdate = indate
    elif sign.lower() in ['b', 'begin']:
        outdate = indate.replace(day=1).replace(month=1)
    elif sign.lower() in ['e', 'end']:
        outdate = indate.replace(day=1).replace(month=12).replace(day=31)
    else:
        raise ValueError('Invalid sign')
    return outdate


def _month(indate: datetime, sign: Literal['s', 'b', 'e'] = 's') -> datetime:
    if sign.lower() in ['s', 'same']:
        outdate = indate
    elif sign.lower() in ['b', 'begin']:
        outdate = indate.replace(day=1)
    elif sign.lower() in ['e', 'end']:
        outdate = indate + relativedelta(months=1)
        outdate = outdate.replace(day=1) + relativedelta(days=-1)
    else:
        raise ValueError('Invalid sign')
    return outdate


def _quarter(indate: datetime, sign: Literal['s', 'b', 'e'] = 's') -> datetime:
    if sign.lower() in ['s', 'same']:
        outdate = indate
    elif sign.lower() in ['b', 'begin']:
        outdate = indate.replace(day=1).replace(month=3 * ((indate.month - 1) // 3) + 1)
    elif sign.lower() in ['e', 'end']:
        outdate = indate.replace(day=1).replace(month=3 * ((indate.month - 1) // 3) + 3)
    else:
        raise ValueError('Invalid sign')
    return _month(outdate, 'e')


def _week(indate: datetime, sign: Literal['s', 'b', 'e'] = 's') -> datetime:
    if sign.lower() in ['s', 'same']:
        outdate = indate
    elif sign.lower() in ['b', 'begin']:
        outdate = indate - relativedelta(days=datetime.weekday(indate))
    elif sign.lower() in ['e', 'end']:
        outdate = indate + relativedelta(days=6 - datetime.weekday(indate))
    else:
        raise ValueError('Invalid sign')
    return outdate


def intnx(fromdate: str = TODAY,
          interval: Literal['y', 'q', 'm', 'w', 'd'] = 'm',
          n: int = -1,
          sign: Literal['s', 'b', 'e'] = 's') -> str:
    """ 计算从from开始经过n个in间隔后的SAS日期 """
    fromdate = datetime.strptime(fromdate, '%Y-%m-%d')

    # 按照月份进行处理
    if interval.lower() in ['y', 'year']:
        outdate = fromdate + relativedelta(years=n)
        outdate = _year(outdate, sign)
    elif interval.lower() in ['q', 'quarter']:
        outdate = fromdate + relativedelta(months=3 * n)
        outdate = _quarter(outdate, sign)
    elif interval.lower() in ['m', 'month']:
        outdate = fromdate + relativedelta(months=n)
        outdate = _month(outdate, sign)
    elif interval.lower() in ['w', 'week']:
        outdate = fromdate + relativedelta(weeks=n)
        outdate = _week(outdate, sign)
    elif interval.lower() in ['d', 'day']:
        outdate = fromdate + relativedelta(days=n)
    else:
        raise ValueError('Invalid interval')
    return outdate.strftime('%Y-%m-%d')


def intck(fromdate=TODAY, todate=TODAY, interval='m'):
    """ 计算从from开始到todate之间的间隔 """
    fromdate = datetime.strptime(fromdate, '%Y-%m-%d')
    todate = datetime.strptime(todate, '%Y-%m-%d')
    if interval.lower() in ['d', 'day']:
        timediff = todate - fromdate
        return timediff.days
    elif interval.lower() in ['m', 'month']:
        return 12 * (todate.year - fromdate.year) + (todate.month - fromdate.month)
    elif interval.lower() in ['y', 'year']:
        return todate.year - fromdate.year


if __name__ == '__main__':
    print(intck('2012-01-13', '2022-02-13', 'd'))
    print(intnx('2012-02-18', 'm', -12))
    print(intnx('2012-02-18', 'q', n=-1, sign='b'))
    print(intnx('2012-02-18', 'q', n=-1, sign='e'))
    print(intnx('2012-04-18', 'q', n=-2, sign='b'))
    print(intnx('2020-05-31', 'q', n=0, sign='e'))
